%matplotlib widget
from utils.traffic_model import SignalController

# 碁盤目状都市の大きさを指定。都市は city_size*city_size の格子状のマップで表現される
city_size = 3

# 与えられた都市の大きさに対応した信号機制御パラメータ生成クラスのインスタンス化
signal_controller = SignalController(city_size)

# 整数インデックスとパラメータ値の対応表の表示
signal_controller.show_index_table()

import numpy as np

# 整数乱数により、0 から 19 までの整数インデックスを 3 * city_size * city_size 個生成
index_list = np.random.randint(0, 20, 3 * city_size * city_size)

# 整数インデックスを対応するパラメータ値に変換（デコード）し、制御パラメータを生成
signal_info = signal_controller.decode(index_list=index_list, disp_flag=True)

from utils.traffic_model import TrafficSimulator2Malls

# シミュレータのインスタンス化
traffic_sim = TrafficSimulator2Malls(
    signal_info=signal_info, num_cars=20, ratio_mall=0.5, seed=0
)

# シミュレーションの実行（アニメーション表示）
traffic_sim.run(steps=500, animation=True)

# シミュレーションの実行（gif アニメーションの出力）
# ※下記引数の設定でおよそ1分程度の処理時間
# ※gif出力はローカル環境でお試しください。
# traffic_sim.run(steps=200, gif="./anim.gif")

import os
import torch


def seed_everything(seed=0):
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

from amplify import FixstarsClient
from datetime import timedelta

client = FixstarsClient()
client.parameters.timeout = timedelta(milliseconds=5000)  # 5000 ミリ秒
# client.token = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"  # ローカル環境等で使用する場合は、Amplify AEのアクセストークンを入力してください。

import torch.nn as nn


class TorchFM(nn.Module):
    def __init__(self, d: int, k: int) -> None:
        super().__init__()
        self.V = nn.Parameter(torch.randn(d, k), requires_grad=True)
        self.lin = nn.Linear(d, 1)

    def forward(self, x):
        out_1 = torch.matmul(x, self.V).pow(2).sum(1, keepdim=True)
        out_2 = torch.matmul(x.pow(2), self.V.pow(2)).sum(1, keepdim=True)
        out_inter = 0.5 * (out_1 - out_2)
        out_lin = self.lin(x)
        out = out_inter + out_lin
        return out

from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from typing import Type

import copy


def train(
    X: np.ndarray,
    y: np.ndarray,
    model_class: Type[nn.Module],
    model_params: dict[str, int | float],
    batch_size=1024,
    epochs=3000,
    criterion=nn.MSELoss(),
    optimizer_class=torch.optim.AdamW,
    opt_params={"lr": 1},
    lr_sche_class=None,
    lr_sche_params=None,
):
    X_tensor, y_tensor = (
        torch.from_numpy(X).float(),
        torch.from_numpy(y).float(),
    )
    indices = np.array(range(X.shape[0]))
    indices_train, indices_valid = train_test_split(
        indices, test_size=0.2, random_state=42
    )

    train_set = TensorDataset(X_tensor[indices_train], y_tensor[indices_train])
    valid_set = TensorDataset(X_tensor[indices_valid], y_tensor[indices_valid])
    loaders = {
        "train": DataLoader(train_set, batch_size=batch_size, shuffle=True),
        "valid": DataLoader(valid_set, batch_size=batch_size, shuffle=False),
    }

    model = model_class(**model_params)
    best_model_wts = copy.deepcopy(model.state_dict())
    optimizer = optimizer_class(model.parameters(), **opt_params)
    scheduler = None
    if lr_sche_class is not None:
        scheduler = lr_sche_class(optimizer, **lr_sche_params)
    best_score = 1e18
    for _ in range(epochs):
        losses = {"train": 0.0, "valid": 0.0}

        for phase in ["train", "valid"]:
            if phase == "train":
                model.train()
            else:
                model.eval()

            for batch_x, batch_y in loaders[phase]:
                optimizer.zero_grad()
                out = model(batch_x).T[0]
                loss = criterion(out, batch_y)
                losses[phase] += loss.item() * batch_x.size(0)

                with torch.set_grad_enabled(phase == "train"):
                    if phase == "train":
                        loss.backward()
                        optimizer.step()

            losses[phase] /= len(loaders[phase].dataset)

        with torch.no_grad():
            model.eval()
            if best_score > losses["valid"]:
                best_model_wts = copy.deepcopy(model.state_dict())
                best_score = losses["valid"]
        if scheduler is not None:
            scheduler.step()

    with torch.no_grad():
        model.load_state_dict(best_model_wts)
        model.eval()

    out = model(X_tensor).T[0]
    coef = torch.corrcoef(torch.stack((out, y_tensor)))[0, 1].detach()
    return model, coef.item()

def gen_training_data_flow_onehot(D: int, N0: int, true_func, input_generator):
    X = np.zeros((N0, D))
    y = np.zeros(N0)
    print("")
    print("###################################################")
    print(" Making initial training data")
    print("###################################################")
    for i in range(N0):
        print(f"Generating training data set #{i}.")
        # (0 or 1)の要素からなる入力ベクトルをランダムに決定
        x = input_generator()
        # 既に全く同じ入力ベクトルが与えられている場合、入力値を再生成
        is_identical = True
        while is_identical:
            is_identical = False
            for j in range(i + 1):
                if np.all(x == X[j, :]):
                    x = input_generator()
                    is_identical = True
                    break
        # 目的関数 f(x) で入力値 x を評価し、入出力ペア (x,f(x)) を教師データにコピー
        X[i, :] = x
        y[i] = true_func(x)
        print(f"------------------------")
    return X, y

from amplify import VariableGenerator, solve, one_hot, sum
import matplotlib.pyplot as plt
import sys
import time

# FMQA for non-binary inputs


class FMQA_NB:
    def __init__(self, D, N, N0, k, true_func, client, encoder_decoder) -> None:
        assert N0 < N
        self.D = D
        self.N = N
        self.N0 = N0
        self.k = k
        self.true_func = true_func
        self.client = client
        self.nbins_array = encoder_decoder.nbins_array
        self.encoder = encoder_decoder.encoder
        self.decoder = encoder_decoder.decoder
        self.y = None

    # 教師データに基づいて N-N0 回の FMQA を教師データを追加しながら繰り返し実施するメンバー関数
    def cycle(self, X, y, log: bool = False) -> np.ndarray:
        # one-hot 制約条件の重み係数
        constraint_weight = max([1, int(np.abs(y).max() + 0.5)])
        print("")
        print("###################################################")
        print(f" Starting FMQA cycles... {constraint_weight=}")
        print("###################################################")

        pred_x = X[0]
        pred_y = 1e18
        self.y = y
        i_sta = 0
        for i in range(self.N - self.N0):
            print(f"FMQA Cycle #{i} ")
            start_time = time.perf_counter()
            try:
                x_hat = self.step(X, y, constraint_weight)
            except RuntimeError:
                sys.exit(f"Unknown error, i = {i}")
            # x_hat として既に全く同じ入力が教師データ内に存在する場合、その周辺の値を x_hat とする。
            is_identical = True
            while is_identical:
                is_identical = False
                for j in range(i + self.N0):
                    if np.all(x_hat == X[j, :]):
                        # バイナリベクトルから整数ベクトルへデコードし、inputs へコピー
                        inputs = self.decoder(x_hat)
                        # 周辺の値を取得（inputs は整数ベクトル）
                        inputs += np.random.randint(-1, 2, len(inputs))
                        for i_inp in range(len(inputs)):
                            if inputs[i_inp] < 0:
                                inputs[i_inp] = 0
                            elif inputs[i_inp] > self.nbins_array[i_inp] - 1:
                                inputs[i_inp] = self.nbins_array[i_inp] - 1
                        # 整数ベクトルからバイナリベクトルへエンコードし、x_hat へコピー
                        x_hat = self.encoder(inputs)
                        if log:
                            print(f"{i=}, Identical x is found, {x_hat=}")
                        is_identical = True
                        break
            # hat{x} で目的関数 f() を評価
            y_hat = self.true_func(x_hat)
            # 最適点近傍における入出力ペア [x_hat, y_hat] を教師データに追加
            X = np.vstack((X, x_hat))
            y = np.append(y, y_hat)
            self.y = np.append(self.y, y_hat)
            # 目的関数の評価値が最小値を更新したら、その入出力ペアを [pred_x, pred_y] へコピー
            print(
                f"FMQA Cycle #{i} finished in {time.perf_counter() - start_time:.2f}s",
                end="",
            )
            if pred_y > y_hat:
                pred_y = y_hat
                pred_x = x_hat
                print(f", variable updated, {pred_y=:.2f}")
            else:
                print()
            print(f"------------------------")
        return pred_x

    # 1回のFMQAを実施するメンバー関数
    def step(self, X, y, constraint_weight) -> np.ndarray:
        # FM を機械学習
        start_time = time.perf_counter()
        model, corrcoef = train(
            X,
            y,
            model_class=TorchFM,
            model_params={"d": self.D, "k": self.k},
            batch_size=8,
            epochs=1000,
            criterion=nn.MSELoss(),
            optimizer_class=torch.optim.AdamW,
            # 学習率をエポック数とともに小さくするスケジューラーを設定
            opt_params={"lr": 0.5},
            lr_sche_class=torch.optim.lr_scheduler.StepLR,
            lr_sche_params={"step_size": 50, "gamma": 0.9},
        )
        print(
            f"FM Training: {time.perf_counter() - start_time:.2f}s (corr:{corrcoef:.2f})"
        )
        start_time = time.perf_counter()
        # 学習済みモデルから、FM パラメータの抽出
        v, w, w0 = list(model.parameters())
        v = v.detach().numpy()
        w = w.detach().numpy()[0]
        w0 = w0.detach().numpy()[0]
        # ここから量子アニーリング (QA) を実施
        gen = VariableGenerator()  # 変数ジェネレータを宣言
        q = gen.array("Binary", self.D)  # 決定変数の作成
        cost = self.__FM_as_QUBO(q, w0, w, v)  # FM パラメータから QUBO として FM を定義
        # 一つの変数が一つの値を持つように one-hot 制約
        constraint_list = []
        ista = 0
        iend = 0
        for i in range(len(self.nbins_array)):
            iend += self.nbins_array[i]
            constraint_list.append(one_hot(q[ista:iend]))
            ista = iend
        constraints = sum(constraint_list)
        # 目的関数と制約条件を足し合わし、Amplify のソルバーに受け渡し
        model = cost + constraint_weight * constraints
        result = solve(model, self.client)  # 目的関数を Amplify のソルバーに受け渡し
        print(
            f"QUBO: {time.perf_counter() - start_time:.2f}s (AE execution: {result.execution_time.seconds:.2f}s)"
        )
        if len(result.solutions) == 0:
            raise RuntimeError("No solution was found.")
        q_values = q.evaluate(result.best.values)
        return q_values

    # FM パラメータから QUBO として FM を定義する関数。前定義の TorchFM クラスと同様に、g(x) の関数形通りに数式を記述。
    def __FM_as_QUBO(self, x, w0, w, v):
        lin = w0 + (x.T @ w)
        out_1 = np.array([(x * v[:, i]).sum() ** 2 for i in range(self.k)]).sum()
        # 次式において、x[j] はバイナリ変数なので、x[j] = x[j]^2 であることに注意。
        out_2 = np.array([(x * v[:, i] * v[:, i]).sum() for i in range(self.k)]).sum()
        return lin + (out_1 - out_2) / 2

    # 初期教師データ及び各 FMQA サイクル内で実施した N 回の目的関数評価値の履歴をプロットする関数
    def plot_history(self):
        assert self.y is not None
        fig = plt.figure(figsize=(6, 4))
        plt.plot(
            [i for i in range(self.N0)],
            self.y[: self.N0],
            marker="o",
            linestyle="-",
            color="b",
        )  # 初期教師データ生成時の目的関数評価値
        plt.plot(
            [i for i in range(self.N0, self.N)],
            self.y[self.N0 :],
            marker="o",
            linestyle="-",
            color="r",
        )  # FMQA サイクル時の目的関数評価値
        plt.plot(
            [i for i in range(self.N)],
            [self.y[:i].min() for i in range(1, self.N + 1)],
            linestyle="-",
            color="k",
        )  # 目的関数最小値の更新履歴
        plt.xlabel("i-th evaluation of f(x)", fontsize=18)
        plt.ylabel("f(x)", fontsize=18)
        plt.tick_params(labelsize=18)
        return fig

class EncDecOneHot:
    def __init__(self, D: int, nbins_array):
        self.D = D
        self.nbins_array = nbins_array

    def gen_random_input(self):
        ista = 0
        iend = 0
        x = np.zeros(self.D, int)
        for i in range(len(self.nbins_array)):
            iend += self.nbins_array[i]
            idx = np.random.randint(ista, iend, 1)[0]
            x[idx] = 1
            ista = iend
        return x

    def encoder(self, inputs):
        if len(inputs) != len(self.nbins_array):
            raise RuntimeError("inputs should be the same length as nbins_array!")
        x = np.zeros(self.D, int)
        i_offset = 0
        for i in range(len(inputs)):
            x[inputs[i] + i_offset] = 1
            i_offset += self.nbins_array[i]
        return x

    def decoder(self, x):
        bit_locs = np.where(x == 1)[0]
        if len(bit_locs) != len(self.nbins_array):
            raise RuntimeError("bit_locs should be the same length as nbins_array!")
        i_offset = 0
        for i in range(1, len(bit_locs)):
            i_offset += self.nbins_array[i - 1]
            bit_locs[i] -= i_offset
        return bit_locs

import time

# 乱数シード値を初期化
seed_everything()

# 碁盤目状都市の大きさを指定。city_size*city_size の格子マップ
city_size = 4

# 信号設定生成クラスのインスタンス化
signal_controller = SignalController(city_size)

# 各変数の one-hot エンコーディングの bin サイズの1次元配列
nbins_array = np.array(
    [
        [
            signal_controller.nbins_red_duration,
            signal_controller.nbins_green_duration,
            signal_controller.nbins_phase,
        ]
        for i in range(city_size**2)
    ]
).reshape(-1)

# バイナリ変数の全体のサイズ（QUBO が取り扱うバイナリ決定変数のサイズ）
D = nbins_array.sum()
print(f"Number of binary decision variables: {D}")

# One-hot エンコーダー・デコーダークラス（EncDecOneHot）のインスタンス化
enc_dec_one_hot = EncDecOneHot(D, nbins_array)

# 目的関数評価回数のカウンター
n_eval = 0

num_cars = 100
ratio_mall = 0.5


# 目的関数（信号機制御パラメータの取得 → 交通シミュレーション実施 → 目的関数値の計算）
def obj_func(x):
    global n_eval
    # 求解結果（バイナリ）から、信号機制御パラメータ（赤の期間、緑の期間、位相）を表現する整数インデックスに変換する。
    index_list = enc_dec_one_hot.decoder(x)

    # さらに、signal_controller.decode メソッドにより、整数インデックスから信号機制御パラメータ（実数）に変換する。
    signal_info = signal_controller.decode(index_list)

    # 同じ信号機制御で、n_cases 回の異なるシミュレーションを実施。
    n_cases = 5
    costs = []
    print("Cost convergence: ", end="")
    start_time = time.perf_counter()
    for i in range(n_cases):
        # TrafficSimulator2Malls のインスタンス化（初期条件を決定するシード値は乱数）
        traffic_sim = TrafficSimulator2Malls(
            signal_info=signal_info,
            num_cars=num_cars,
            ratio_mall=ratio_mall,
            seed=np.random.randint(0, 100),
        )
        # シミュレーションの実行
        traffic_sim.run(steps=20000)
        # 各シミュレーションにおける全自動車の平均速度の負値を取得し、costs に追加
        costs.append(-traffic_sim.get_latest_mean())
        print(f"{sum(costs) / (i + 1):.2f} ", end="")

    # n_cases 回のシミュレーションの平均出力値を最終的な目的関数値とする。
    cost = sum(costs) / n_cases
    print(
        f"\nEvaluation #{n_eval} finished in {time.perf_counter() - start_time:.2f}s. cost:{cost:.2f}"
    )
    n_eval += 1

    return cost


N = 3  # 関数を評価できる回数。『3.4. 本サンプルコードによる FMQA 実行例』では 100 と設定
N0 = 2  # 初期教師データのサンプル数。『3.4. 本サンプルコードによる FMQA 実行例』では 10 と設定
k = 20  # FMにおけるベクトルの次元（ハイパーパラメータ）

# 時間計測用開始時刻
start_time = time.perf_counter()

# 初期教師データの生成
X, y = gen_training_data_flow_onehot(D, N0, obj_func, enc_dec_one_hot.gen_random_input)

# FMQA のインスタンス化
fmqa_solver = FMQA_NB(
    D=D,
    N=N,
    N0=N0,
    k=k,
    true_func=obj_func,
    client=client,
    encoder_decoder=enc_dec_one_hot,
)

# FMQA サイクルの実行
pred_x = fmqa_solver.cycle(X, y)

# 最適化結果の出力
print("###################################################")
print(" Optimal input values and corresponding output")
print("###################################################")
print(f"pred x: {pred_x}")
print(f"pred value: {obj_func(pred_x):.2f}")
print(f"Elapsed time: {time.perf_counter() - start_time:.2f}s")

fig = fmqa_solver.plot_history()

# 最適化の際のシミュレーションと同じ条件設定
city_size = 4
num_cars = 100
ratio_mall = 0.5
signal_controller = SignalController(city_size)

# 最適化結果 pred_x から整数インデックスを構築
index_list = enc_dec_one_hot.decoder(pred_x)
print(f"{index_list=}")

# 整数インデックスを信号機制御パラメータに変換
signal_info = signal_controller.decode(index_list)

# シミュレータのインスタンス化
traffic_sim = TrafficSimulator2Malls(
    signal_info=signal_info, num_cars=num_cars, ratio_mall=ratio_mall
)

# シミュレーションの実施
traffic_sim.run(2000, animation=True, interval=100)

  Number of binary decision variables: 960

  ###################################################
  Making initial training data
  ###################################################
  Generating training data set #0.
  Cost convergence: -2.40 -3.41 -3.32 -3.53 -3.95 
  Evaluation #0 finished in 52.88s. cost:-3.95
  ------------------------
  Generating training data set #1.
  Cost convergence: -6.56 -6.06 -6.13 -6.00 -5.50 
  Evaluation #1 finished in 50.06s. cost:-5.50
  ------------------------
  Generating training data set #2.
  Cost convergence: -6.19 -6.58 -6.59 -6.56 -6.49 
  Evaluation #2 finished in 47.57s. cost:-6.49
  ------------------------
  Generating training data set #3.
  Cost convergence: -4.18 -4.01 -4.07 -4.10 -4.13 
  Evaluation #3 finished in 53.65s. cost:-4.13
  ------------------------
  Generating training data set #4.
  Cost convergence: -3.79 -4.32 -4.19 -3.73 -3.90 
  Evaluation #4 finished in 55.38s. cost:-3.90
  ------------------------
  Generating training data set #5.
  Cost convergence: -1.68 -3.16 -3.20 -3.22 -3.06 
  Evaluation #5 finished in 56.32s. cost:-3.06
  ------------------------
  Generating training data set #6.
  Cost convergence: -6.94 -5.61 -5.75 -5.71 -5.90 
  Evaluation #6 finished in 50.48s. cost:-5.90
  ------------------------
  Generating training data set #7.
  Cost convergence: -5.38 -6.08 -5.45 -5.38 -5.43 
  Evaluation #7 finished in 51.79s. cost:-5.43
  ------------------------
  Generating training data set #8.
  Cost convergence: -3.43 -3.91 -4.12 -4.55 -4.85 
  Evaluation #8 finished in 50.39s. cost:-4.85
  ------------------------
  Generating training data set #9.
  Cost convergence: -3.77 -4.18 -4.31 -4.37 -4.08 
  Evaluation #9 finished in 54.33s. cost:-4.08
  ------------------------

ブラックボックス最適化による信号制御最適化¶

1. 都市渋滞の MAS¶

1.1. MASとは¶

1.2. 交通渋滞モデル¶

1.3. MAS 交通シミュレータ¶

2. FMQA のプログラム実装（整数入力値）¶

2.1. 乱数の初期化¶

2.2. クライアントの設定¶

2.3. PyTorch による FM の実装¶

2.4. 初期教師データの作成¶

2.5. FMQA サイクルの実行クラス¶

2.6. One-hot エンコーダー・デコーダーの実装¶

3. ブラックボックス最適化による信号機制御の最適化¶

3.1. 最適化の実行¶

3.2. 最適化過程における目的関数値の推移¶

3.3. 最適化された信号機設定での交通シミュレーション¶

3.4. 本サンプルコードによる実行例¶

4. まとめ及び実運用の指針¶